"""
Dataset Generator
author: RuanQuanyuan@2022/4/21

Data Format
'name of pic' 'label'
instance:
001.jpg 0
002.jpg 1
...
"""
import logging
import os
import cv2
import numpy as np


class DatasetGenerator:
    def __init__(self, data_path: str, label_path: str, resize=None):
        self.data = []
        self.label = []

        self.resize = resize

        with open(label_path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                data = line.strip().split(' ')
                self.data.append(data_path + data[0])
                self.label.append(int(data[1]))

    def __getitem__(self, item):
        assert os.path.exists(self.data[item]), "image:{} don't exist.".format(self.data[item])

        img = cv2.imread(self.data[item])

        img.astype(np.float32)

        if self.resize:
            img = cv2.resize(img, (self.resize, self.resize)) / 255  # 放缩归一化
        else:
            img = img / 255  # 归一化

        # add additional Gaussian Noise
        img += np.random.normal(loc=0.0, scale=0.001, size=img.shape,)

        label = self.label[item]

        return img, label

    def __len__(self):
        return len(self.label)
